import jwt
import json
import asyncio
import aiohttp
from functools import wraps
from flask import request, jsonify, g
from datetime import datetime, timedelta
from app.models import User, APIKey, DoesNotExist

import logging, sys
logging.basicConfig(
    level=logging.INFO,
    format='[UTILS222] %(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler(sys.stdout)]   # 显式指向控制台
)
logger = logging.getLogger('ollama')

# JWT 工具函数
def generate_jwt(user_id, app):
    """生成JWT令牌"""
    payload = {
        'user_id': user_id,
        'exp': datetime.utcnow() + timedelta(seconds=app.config['JWT_EXPIRE_SEC'])
    }
    return jwt.encode(payload, app.config['SECRET_KEY'], algorithm='HS256')

def verify_jwt(token, app):
    """验证JWT令牌"""
    try:
        payload = jwt.decode(token, app.config['SECRET_KEY'], algorithms=['HS256'])
        return payload['user_id']
    except (jwt.ExpiredSignatureError, jwt.InvalidTokenError):
        return None

# 装饰器：登录验证
def login_required(f):
    """登录验证装饰器，支持异步视图"""
    @wraps(f)
    async def wrapper(*args, **kwargs):
        auth_header = request.headers.get('Authorization')
        if not auth_header or not auth_header.startswith('Bearer '):
            return jsonify({'error': '未提供Bearer令牌'}), 401
        
        user_id = verify_jwt(auth_header.split(' ')[1], g.app)
        if not user_id:
            return jsonify({'error': '令牌无效或已过期'}), 401
        
        try:
            g.user = User.get(User.id == user_id)
        except DoesNotExist:
            return jsonify({'error': '用户不存在'}), 401
        
        # 处理异步和同步视图
        if asyncio.iscoroutinefunction(f):
            return await f(*args, **kwargs)
        return f(*args, **kwargs)
    return wrapper

# 装饰器：管理员权限验证
def admin_required(f):
    """管理员权限验证装饰器"""
    @wraps(f)
    @login_required
    async def wrapper(*args, **kwargs):
        if not g.user.is_admin:
            return jsonify({'error': '需要管理员权限'}), 403
        
        if asyncio.iscoroutinefunction(f):
            return await f(*args, **kwargs)
        return f(*args, **kwargs)
    return wrapper

# 装饰器：API Key验证
def api_key_required(f):
    """API Key验证装饰器"""
    @wraps(f)
    @login_required
    async def wrapper(*args, **kwargs):
        api_key = request.headers.get('X-API-Key')
        if not api_key:
            return jsonify({'error': '未提供X-API-Key头'}), 401
        
        try:
            g.api_key = APIKey.get(
                (APIKey.key == api_key) &
                (APIKey.user == g.user) &
                (APIKey.is_active == True) &
                (APIKey.is_revoked == False)
            )
        except DoesNotExist:
            return jsonify({'error': 'API Key无效/未激活/已注销'}), 401
        
        if asyncio.iscoroutinefunction(f):
            return await f(*args, **kwargs)
        return f(*args, **kwargs)
    return wrapper


# 装饰器：检查用户是否有激活的 Key验证
def auto_api_key(f):
    """自动校验当前用户是否存在【已激活】的 APIKey"""
    @wraps(f)
    @login_required
    async def wrapper(*args, **kwargs):
        try:
            key = APIKey.get(
                (APIKey.user == g.user) &
                (APIKey.is_active == True) &
                (APIKey.is_revoked == False)
            )
            g.api_key = key          # 供后续流式调用使用
        except DoesNotExist:
            return jsonify({'error': '您没有已激活的 APIKey，请先申请并等待管理员激活'}), 403

        if asyncio.iscoroutinefunction(f):
            return await f(*args, **kwargs)
        return f(*args, **kwargs)
    return wrapper

# AI调用工具
async def async_ai_stream(prompt: str, history: list, server_url: str, model_name: str):
    """异步流式调用AI接口（Ollama兼容格式）"""

    print(history)
    print(server_url)
    print(model_name)
    print(prompt)

        # ---------- 1. 连通性探测 ----------
    try:
        # 拼一个最轻量的 /api/tags 请求，只要 200 就认为 alive
        probe_url = f"{server_url}/api/tags"
        async with aiohttp.ClientSession() as probe_session:
            async with probe_session.get(probe_url, timeout=3) as probe_resp:
                logger.info('[PROBE] %s -> %s', probe_url, probe_resp.status)
                if probe_resp.status != 200:
                    yield f"data: {json.dumps({'error': f'Ollama 连通失败：HTTP {probe_resp.status}'})}\n\n"
                    yield "data: [DONE]\n\n"
                    return
    except Exception as e:
        logger.error('[PROBE] %s 不可达：%s', server_url, e)
        yield f"data: {json.dumps({'error': f'Ollama 不可达：{str(e)}'})}\n\n"
        yield "data: [DONE]\n\n"
        return

    # ---------- 2. 打印真正要发的聊天数据 ----------
    payload = {
        'model': model_name,
        'messages': history + [{'role': 'user', 'content': prompt}],
        'stream': True,
        'options': {'temperature': 0.7}
    }
    logger.info('[CHAT] POST %s/api/chat | model=%s | history_len=%d | prompt=%.50s...',
                server_url, model_name, len(history), prompt)

    # ---------- 3. 正式流式请求 ----------

    timeout = aiohttp.ClientTimeout(total=g.app.config['AI_TIMEOUT'])
    async with aiohttp.ClientSession(timeout=timeout) as session:
        try:
            async with session.post(
                url=f"{server_url}/api/chat",
                headers={'Content-Type': 'application/json'},
                json={
                    'model': model_name,
                    'messages': history + [{'role': 'user', 'content': prompt}],
                    'stream': True,
                    'options': {'temperature': 0.7}
                }
            ) as resp:
                resp.raise_for_status()
                
                async for line in resp.content:
                    line = line.strip().decode('utf-8')
                    if not line:
                        continue
                    
                    try:
                        data = json.loads(line)
                        if 'message' in data and 'content' in data['message']:
                            content = data['message']['content']
                            if content:
                                yield f"data: {json.dumps({'content': content})}\n\n"
                    except json.JSONDecodeError:
                        continue
        
        except Exception as e:
            yield f"data: {json.dumps({'error': f'AI调用失败：{str(e)}'})}\n\n"
        finally:
            yield "data: [DONE]\n\n"
